// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.analysis;

import org.apache.doris.nereids.exceptions.AnalysisException;

import java.util.Map;

public class AnnIndexPropertiesChecker {
    public static void checkProperties(Map<String, String> properties) {
        String type = null;
        String metric = null;
        String dim = null;
        String quantizer = null;
        int dimension = 0;
        int numSubQuantizers = 0;
        int nlist = 0;
        for (String key : properties.keySet()) {
            switch (key) {
                case "index_type":
                    type = properties.get(key);
                    if (!type.equals("hnsw") && !type.equals("ivf")) {
                        throw new AnalysisException("only support ann index with type hnsw or ivf, got: " + type);
                    }
                    break;
                case "metric_type":
                    metric = properties.get(key);
                    if (!metric.equals("l2_distance") && !metric.equals("inner_product")) {
                        throw new AnalysisException(
                                "only support ann index with metric l2_distance or inner_product, got: " + metric);
                    }
                    break;
                case "dim":
                    dim = properties.get(key);
                    try {
                        dimension = Integer.parseInt(dim);
                        if (dimension <= 0) {
                            throw new AnalysisException("dim of ann index must be a positive integer, got: " + dim);
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException("dim of ann index must be a positive integer, got: " + dim);
                    }
                    break;
                case "max_degree":
                    String maxDegree = properties.get(key);
                    try {
                        int degree = Integer.parseInt(maxDegree);
                        if (degree <= 0) {
                            throw new AnalysisException(
                                    "max_degree of ann index must be a positive integer, got: " + maxDegree);
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException(
                                "max_degree of ann index must be a positive integer, got: " + maxDegree);
                    }
                    break;
                case "ef_construction":
                    String efConstruction = properties.get(key);
                    try {
                        int ef = Integer.parseInt(efConstruction);
                        if (ef <= 0) {
                            throw new AnalysisException(
                                    "ef_construction of ann index must be a positive integer, got: " + efConstruction);
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException(
                                "ef_construction of ann index must be a positive integer, got: " + efConstruction);
                    }
                    break;
                case "quantizer":
                    quantizer = properties.get(key);
                    if (!quantizer.equals("flat")
                            && !quantizer.equals("sq4")
                            && !quantizer.equals("sq8")
                            && !quantizer.equals("pq")) {
                        throw new AnalysisException(
                                "only support ann index with quantizer flat, sq4, sq8 or pq, got: " + quantizer);
                    }
                    break;
                case "pq_m":
                    String pqM = properties.get(key);
                    try {
                        numSubQuantizers = Integer.parseInt(pqM);
                        if (numSubQuantizers <= 0) {
                            throw new AnalysisException(
                                    "pq_m of ann index must be a positive integer, got: " + numSubQuantizers);
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException(
                                "pq_m of ann index must be a positive integer, got: " + pqM);
                    }
                    break;
                case "pq_nbits":
                    String pqNbits = properties.get(key);
                    try {
                        int numBits = Integer.parseInt(pqNbits);
                        if (numBits <= 0) {
                            throw new AnalysisException(
                                    "pq_nbits of ann index must be a positive integer, got: " + numBits);
                        }
                        if (numBits > 24) {
                            throw new AnalysisException("pq_nbits larger than 24 is not practical");
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException(
                                "pq_nbits of ann index must be a positive integer, got: " + pqNbits);
                    }
                    break;
                case "nlist":
                    String nlistStr = properties.get(key);
                    try {
                        nlist = Integer.parseInt(nlistStr);
                        if (nlist <= 0) {
                            throw new AnalysisException(
                                "nlist of ann index must be a positive integer, got: " + nlistStr);
                        }
                    } catch (NumberFormatException e) {
                        throw new AnalysisException("nlist of ann index must be a positive integer, got: " + nlistStr);
                    }
                    break;
                default:
                    throw new AnalysisException("unknown ann index property: " + key);
            }
        }

        if (type != null && type.equals("ivf")) {
            if (nlist == 0) {
                throw new AnalysisException("nlist of ann index must be specified for ivf type");
            }
        }

        if (type == null) {
            throw new AnalysisException("index_type of ann index be specified.");
        }
        if (metric == null) {
            throw new AnalysisException("metric_type of ann index must be specified.");
        }
        if (dim == null) {
            throw new AnalysisException("dim of ann index must be specified");
        }

        if (quantizer != null && quantizer.equals("pq")) {
            if (dimension == 0 || numSubQuantizers == 0) {
                throw new AnalysisException(
                    "The dimension of the vector (dim) or the number of subquantizers (pq_m) cannot be zero");
            }
            if (dimension % numSubQuantizers != 0) {
                throw new AnalysisException("The dimension of the vector (dim) should be a multiple of the number of "
                        + "subquantizers (pq_m)");
            }
        }
    }
}
